import java.util.Map;
import java.util.TreeMap;

/**
 * 726. 原子的数量
 * https://leetcode-cn.com/problems/number-of-atoms/
 */
public class Solutions_726 {
    public static void main(String[] args) {
//        String formula = "H2O";  // output: "H2O"
//        解释: 原子的数量是 {'H': 2, 'O': 1}。

//        String formula = "Mg(OH)2";  // output: "H2MgO2"
//        解释: 原子的数量是 {'H': 2, 'Mg': 1, 'O': 2}。

        String formula = "K4(ON(SO3)2)2";  // output: "K4N2O14S4"
//        解释: 原子的数量是 {'K': 4, 'N': 2, 'O': 14, 'S': 4}。

        String result = countOfAtoms(formula);
        System.out.println(result);
    }

    private static int i = 0;

    public static String countOfAtoms(String formula) {
        i = 0;
        char[] arr = formula.toCharArray();
        Map<String, Integer> map = recursion(arr);
        StringBuilder sb = new StringBuilder();
        // 不需要排序，treeMap 中的键，本身就是有序的
        // 遍历 map
        for (Map.Entry<String, Integer> entry : map.entrySet()) {
            sb.append(entry.getKey());
            if (entry.getValue() > 1) {
                sb.append(entry.getValue());
            }
        }
        return sb.toString();
    }

    public static Map<String, Integer> recursion(char[] arr) {
        Map<String, Integer> map = new TreeMap<>();
        StringBuilder sb = new StringBuilder();
        int counts = 0, len = arr.length;
        while (i < len && arr[i] != ')') {
            if (Character.isDigit(arr[i])) {
                // 数字
                counts = counts * 10 + (arr[i ++] - '0');
            } else if (Character.isLetter(arr[i])) {
                if (sb.length() > 0) {
                    // NO 时，N 和 O 都是原子，那么遍历 O 时，需要将 N 记录到 map 中
                    int add = counts < 1 ? 1 : counts;
                    map.put(sb.toString(), map.getOrDefault(sb.toString(), 0) + add);
                    sb = new StringBuilder();
                    counts = 0;
                }
                sb.append(arr[i++]);
                // 紧随其后的小写字母一定是一个原子，那么也添加到 sb 中
                while (i < len && Character.isLowerCase(arr[i])) {
                    sb.append(arr[i]);
                    i++;
                }
            } else if (arr[i] == '(') {
                // 递归分解得到子问题
                i ++;
                Map<String, Integer> ans = recursion(arr);
                // 递归返回的结果，加入到 map 中
                for (Map.Entry<String, Integer> entry : ans.entrySet()) {
                    String key = entry.getKey();
                    int count = entry.getValue();
                    map.put(key, map.getOrDefault(key, 0) + count);
                }
            }
        }
        if (sb.length() > 0) {
            // 最后一位为大写字母或者数字的情况，最终需要添加到 map 中
            int add = counts < 1 ? 1 : counts;
            map.put(sb.toString(), map.getOrDefault(sb.toString(), 0) + add);
        }
        ++ i;
        int multi = 0;
        // ')' 后的数字，即是乘积
        while (i < len && Character.isDigit(arr[i])) {
            multi = multi * 10 + (arr[i ++] - '0');
        }
        if (multi > 1) {
            // map 中的全部键值，都乘 multi 次
            for (Map.Entry<String, Integer> entry : map.entrySet()) {
                map.put(entry.getKey(), entry.getValue() * multi);
            }
        }
        return map;
    }
}
